﻿using Assets.Scripts.Bvh.Common;
using Assets.Scripts.Bvh.Gpu;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using UnityEngine;

namespace Assets.Scripts.Bvh.Cpu
{
    public class RadixTreeBuilder
    {
        public List<RadixTreeNode> leafNodes;

        public RadixTreeBuilder(List<RadixTreeNode> leafNodes)
        {
            this.leafNodes = leafNodes;
        }

        public void Sort()
        {
            var dataCount = leafNodes.Count;
            for (int i = 0; i < dataCount; i++)
            {
                for (int j = 0; j < dataCount; j++)
                {
                    if (leafNodes[i].mortonCode > leafNodes[j].mortonCode)
                    {
                        var temp = leafNodes[i];
                        leafNodes[i] = leafNodes[j];
                        leafNodes[j] = temp;
                    }
                }
            }
        }

        //节点莫顿码的快速排序算法，并给节点赋予nodeIndex值
        public void QuickSort()
        {
            CpuSorter.QuickSort(leafNodes, 0, leafNodes.Count - 1);

            var node = new RadixTreeNode();
            for (int i = 0; i < leafNodes.Count; i++)
            {
                node = leafNodes[i];
                node.nodeIndex = i;
                leafNodes[i] = node;
            }
        }

        public RadixTreeNode[] BuildRadixTree()
        {
            var count = leafNodes.Count;
            var interiorNodes = new RadixTreeNode[count - 1];
            for (int i = 0; i < interiorNodes.Length; i++)
            {
                var node = RadixTreeNode.Default;
                node.nodeIndex = i;
                interiorNodes[i] = node;
            }

            for (int i = 0; i < interiorNodes.Length; i++)
            {
                var simRight = Sigma(i, i + 1);
                int simLeft = Sigma(i, i - 1);
                var direction = (int)Sign(simRight - simLeft);
                var sigmaMin = Sigma(i, i - direction);
                var l_max = 2;
                while (Sigma(i, i + l_max * direction) > sigmaMin)
                {
                    l_max = l_max * 2;
                }
                var l = 0;
                for (int t = l_max; t >= 1; t /= 2)
                {
                    if (Sigma(i, i + (l + t) * direction) > sigmaMin)
                    {
                        l = l + t;
                    }
                }
                var j = i + l * direction;
                var sigmaNode = Sigma(i, j);
                var s = 0;
                var init = Mathf.CeilToInt((float)l / 2.0f);
                for (int t = init; t >= 1;)
                {
                    if (Sigma(i, i + (s + t) * direction) > sigmaNode)
                    {
                        s = s + t;
                    }

                    if (t == 1)
                    {
                        break;
                    }

                    t = Mathf.CeilToInt((float)t / 2.0f);
                }
                var gamma = i + s * direction + Mathf.Min(direction, 0);
                RadixTreeNode left;
                RadixTreeNode right;
                if (Mathf.Min(i, j) == gamma)
                {
                    left = leafNodes[gamma];

                    //这里为了区分index是内部节点还是叶子节点，做了一个偏移
                    left.nodeIndex += count - 1;
                    left.parentIndex = i;

                    leafNodes[gamma] = left;
                }
                else
                {
                    left = interiorNodes[gamma];
                    interiorNodes[gamma].parentIndex = i;
                }

                if (Mathf.Max(i, j) == gamma + 1)
                {
                    right = leafNodes[gamma + 1];

                    //这里为了区分index是内部节点还是叶子节点，做了一个偏移
                    right.nodeIndex += count - 1;
                    right.parentIndex = i;

                    leafNodes[gamma + 1] = right;
                }
                else
                {
                    right = interiorNodes[gamma + 1];
                    interiorNodes[gamma + 1].parentIndex = i;
                }

                interiorNodes[i].nodeIndex = i;
                interiorNodes[i].leftNodeIndex = left.nodeIndex;
                interiorNodes[i].rightNodeIndex = right.nodeIndex;
            }

            return interiorNodes;
        }

        private int ComputeLeadingZeros(int value)
        {
            var count = 0;
            for (int i = 31; i >= 0; i--)
            {
                if ((value >> i & 0b01) == 0)
                {
                    count++;
                }
                else
                {
                    return count;
                }
            }

            return count;
        }

        private int Sigma(int i, int j)
        {
            if (j < 0 || j > leafNodes.Count - 1) return -1;
            var a = leafNodes[i].mortonCode;
            var b = leafNodes[j].mortonCode;
            var xor = a ^ b;
            var count = ComputeLeadingZeros(xor);
            return count;
        }

        private float Sign(float value)
        {
            return value / Mathf.Abs(value);
        }

        //合并内部节点和叶子节点得到完成的LBvh
        public RadixTreeNode[] CollectNodes(RadixTreeNode[] interiorNodes)
        {
            var LBvh = new RadixTreeNode[interiorNodes.Length + leafNodes.Count];
            int offset = interiorNodes.Length;
            for (int i = 0; i < LBvh.Length; i++)
            {
                if(i< interiorNodes.Length)
                {
                    LBvh[i] = interiorNodes[i];
                }
                else
                {
                    int leafIndex = i - offset;
                    LBvh[i] = leafNodes[leafIndex];
                }
            }
            return LBvh;
        }

        //计算Bvh树的深度
        public int GetDepth(RadixTreeNode[] LBvh, int nodeIndex)
        {
            var node = LBvh[nodeIndex];
            // nodeIndex = -1 说明节点不存在
            if (nodeIndex == -1)
            {
                return 0;
            }
            if (node.leftNodeIndex == -1 && node.rightNodeIndex == -1)
            {
                // 说明只有根节点
                return 1;
            }

            int leftDepth = GetDepth(LBvh, node.leftNodeIndex);
            int rightDepth = GetDepth(LBvh, node.rightNodeIndex);
            return Mathf.Max(leftDepth, rightDepth) + 1;
        }
    }
}
